import socket
import random
import pickle

# 生成私钥的函数
def generate_private_key(p):
    return random.randint(2, p - 2)

# 计算公钥
def generate_public_key(private_key, g, p):
    return pow(g, private_key, p)

# 计算共享密钥
def generate_shared_secret(private_key, other_public_key, p):
    return pow(other_public_key, private_key, p)

# 接收方的实现
def receiver():
    # 公共素数 p 和基数 g
    p = 23  # 公共素数 p
    g = 5   # 基数 g

    # 生成接收方的私钥和公钥
    receiver_private_key = generate_private_key(p)
    receiver_public_key = generate_public_key(receiver_private_key, g, p)

    print(f"A的私钥是: {receiver_private_key}")
    print(f"A的公钥是: {receiver_public_key}")

    # 创建并绑定 socket
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('localhost', 5000))
        s.listen(1)
        print("Receiver is waiting for connection...")

        # 接受发送方的连接
        conn, addr = s.accept()
        with conn:
            print('Connected by', addr)

            # 接收发送方的公钥
            sender_public_key = pickle.loads(conn.recv(1024))
            print(f"Sender's public key: {sender_public_key}")

            # 发送接收方的公钥给发送方
            conn.send(pickle.dumps(receiver_public_key))

            # 计算共享密钥
            shared_secret = generate_shared_secret(receiver_private_key, sender_public_key, p)
            print(f"Receiver's shared secret: {shared_secret}")
            key = shared_secret.to_bytes((shared_secret.bit_length() + 7) // 8, 'big')[:16]
            with open('shared_secret_A_path.txt', 'wb') as shared_secret_B_file:
                shared_secret_B_file.write(key)


# 启动接收方
if __name__ == "__main__":
    receiver()
